from fastapi import FastAPI
from pydantic import BaseModel
from typing import List, Dict
import requests
import uvicorn

# 创建 FastAPI 应用
app = FastAPI()

# Wikidata SPARQL 查询端点与请求头
WIKIDATA_ENDPOINT = "https://query.wikidata.org/sparql"
HEADERS = {"Accept": "application/sparql-results+json"}

# 定义请求体模型
class LabelRequest(BaseModel):
    label: str
    lang: str = "en"

class QIDRequest(BaseModel):
    qid: str
    lang: str = "en"

class PIDRequest(BaseModel):
    pid: str
    lang: str = "en"

class HeadRelationRequest(BaseModel):
    qid: str
    pid: str
    lang: str = "en"

class HeadRelationValueRequest(BaseModel):
    qid: str
    pid: str

class MIDRequest(BaseModel):
    mid: str

# 通用的 SPARQL 查询执行函数
def run_sparql(query: str):
    response = requests.get(WIKIDATA_ENDPOINT, params={"query": query}, headers=HEADERS)
    response.raise_for_status()
    return response.json()["results"]["bindings"]

# 将标签名转换为实体 QID
@app.post("/label2qid", response_model=List[str])
def label2qid(request: LabelRequest):
    query = f'''
    SELECT ?item WHERE {{
      ?item rdfs:label "{request.label}"@{request.lang}.
    }} LIMIT 10
    '''
    results = run_sparql(query)
    return [r["item"]["value"].split("/")[-1] for r in results]

# 将标签名转换为属性 PID
@app.post("/label2pid", response_model=List[str])
def label2pid(request: LabelRequest):
    query = f'''
    SELECT ?property WHERE {{
      ?property a wikibase:Property ;
                rdfs:label "{request.label}"@{request.lang}.
    }} LIMIT 10
    '''
    results = run_sparql(query)
    return [r["property"]["value"].split("/")[-1] for r in results]

# 根据 PID 获取属性的标签（名称）
@app.post("/pid2label", response_model=List[str])
def pid2label(request: PIDRequest):
    query = f'''
    SELECT ?label WHERE {{
      wd:{request.pid} rdfs:label ?label.
      FILTER(LANG(?label) = "{request.lang}")
    }}
    '''
    results = run_sparql(query)
    return [r["label"]["value"] for r in results]

# 根据 QID 获取实体的标签（名称）
@app.post("/qid2label", response_model=List[str])
def qid2label(request: QIDRequest):
    query = f'''
    SELECT ?label WHERE {{
      wd:{request.qid} rdfs:label ?label.
      FILTER(LANG(?label) = "{request.lang}")
    }}
    '''
    results = run_sparql(query)
    return [r["label"]["value"] for r in results]

# 获取某个实体的所有属性及其值（包含属性标签和值标签）
@app.post("/get_all_relations_of_an_entity", response_model=Dict[str, List[str]])
def get_all_relations_of_an_entity(request: QIDRequest):
    query = f'''
    SELECT ?prop ?propLabel ?value ?valueLabel WHERE {{
      wd:{request.qid} ?p ?value.
      ?prop wikibase:directClaim ?p.
      SERVICE wikibase:label {{ bd:serviceParam wikibase:language "{request.lang}". }}
    }}
    '''
    results = run_sparql(query)
    out = {}
    for r in results:
        pid = r["prop"]["value"].split("/")[-1]
        val = r["valueLabel"]["value"] if "valueLabel" in r else r["value"]["value"]
        # 处理有http的情况
        if "http://" in val or "https://" in val:
            continue
        out.setdefault(pid, []).append(val)
    print(len(out))
    return out

# 获取指定实体和属性下的尾实体（QID 与标签）
@app.post("/get_tail_entities_given_head_and_relation", response_model=Dict[str, str])
def get_tail_entities_given_head_and_relation(request: HeadRelationRequest):
    query = f'''
    SELECT ?value ?valueLabel WHERE {{
      wd:{request.qid} wdt:{request.pid} ?value.
      SERVICE wikibase:label {{ bd:serviceParam wikibase:language "{request.lang}". }}
    }}
    '''
    results = run_sparql(query)
    return {r["value"]["value"].split("/")[-1]: r["valueLabel"]["value"] for r in results if "valueLabel" in r}

# 获取指定实体和属性下的值（如出生日期、人口等）
@app.post("/get_tail_values_given_head_and_relation", response_model=List[str])
def get_tail_values_given_head_and_relation(request: HeadRelationValueRequest):
    query = f'''
    SELECT ?value WHERE {{
      wd:{request.qid} wdt:{request.pid} ?value.
    }}
    '''
    results = run_sparql(query)
    return [r["value"]["value"] for r in results]

# 获取外部 ID（如 IMDb ID 等）
@app.post("/get_external_id_given_head_and_relation", response_model=List[str])
def get_external_id_given_head_and_relation(request: HeadRelationValueRequest):
    query = f'''
    SELECT ?value WHERE {{
      wd:{request.qid} wdt:{request.pid} ?value.
    }}
    '''
    results = run_sparql(query)
    return [r["value"]["value"] for r in results]

# 根据 Freebase 中的 M-ID 查找对应 QID（Wikidata）
@app.post("/mid2qid", response_model=List[str])
def mid2qid(request: MIDRequest):
    query = f'''
    SELECT ?item WHERE {{
      ?item wdt:P646 "{request.mid}".
    }} LIMIT 10
    '''
    results = run_sparql(query)
    return [r["item"]["value"].split("/")[-1] for r in results]


if __name__ == '__main__':
    uvicorn.run(app, host="127.0.0.1", port=8088)